#!/usr/bin/env python3
"""
Comprehensive Mapping Table Extractor
Systematically builds the complete 245-byte mapping table
"""

import os
import json
import re
from collections import defaultdict, Counter

def build_comprehensive_mapping():
    """Build comprehensive mapping table using all available information"""
    
    print("=== Building Comprehensive Mapping Table ===")
    
    # Load data
    with open('decoded_result.bin', 'rb') as f:
        encrypted_data = f.read()
    
    # Start with confirmed mappings
    confirmed_mappings = {
        0xa1: ord('m'),     # Confirmed from frequency analysis
        0xf3: ord('o'),     # Confirmed from frequency analysis  
        0x6b: ord('d'),     # Confirmed from frequency analysis
        0x31: ord('e'),     # Confirmed from frequency analysis
        0xb0: ord('l'),     # Confirmed from frequency analysis
        0x69: ord('"'),     # Confirmed from frequency analysis
        0x01: ord(':'),     # Confirmed from frequency analysis
        0x90: ord(' '),     # Confirmed from frequency analysis (space)
        0x52: ord('}'),     # Found in recent analysis
    }
    
    print(f"Starting with {len(confirmed_mappings)} confirmed mappings")
    
    # Analyze patterns to find more mappings
    new_mappings = {}
    
    # Method 1: JSON structure analysis
    json_mappings = analyze_json_structure(encrypted_data, confirmed_mappings)
    new_mappings.update(json_mappings)
    
    # Method 2: Common word analysis
    word_mappings = analyze_common_words(encrypted_data, confirmed_mappings)
    new_mappings.update(word_mappings)
    
    # Method 3: Character frequency analysis
    freq_mappings = analyze_character_frequencies(encrypted_data, confirmed_mappings)
    new_mappings.update(freq_mappings)
    
    # Method 4: Pattern repetition analysis
    pattern_mappings = analyze_pattern_repetitions(encrypted_data, confirmed_mappings)
    new_mappings.update(pattern_mappings)
    
    # Combine all mappings
    all_mappings = {**confirmed_mappings, **new_mappings}
    
    print(f"Total mappings discovered: {len(all_mappings)}")
    
    # Create full mapping table
    mapping_table = create_full_mapping_table(all_mappings)
    
    # Test the mapping table
    test_comprehensive_mapping(mapping_table, encrypted_data)
    
    return mapping_table, all_mappings

def analyze_json_structure(encrypted_data, known_mappings):
    """Analyze JSON structure patterns to find more mappings"""
    
    print("\n--- JSON Structure Analysis ---")
    
    new_mappings = {}
    
    # Common JSON patterns we expect to find
    expected_patterns = [
        (b'{"', 'Object start'),
        (b'"}', 'Object end'), 
        (b',"', 'Field separator'),
        (b'":"', 'Key-value separator'),
        (b'":null', 'Null value'),
        (b'":true', 'Boolean true'),
        (b'":false', 'Boolean false'),
        (b'"questionText"', 'Question field'),
        (b'"sessionId"', 'Session field'),
        (b'"model"', 'Model field'),
    ]
    
    # Use our known mappings to decrypt and look for these patterns
    for pattern_bytes, desc in expected_patterns:
        positions = find_pattern_positions(encrypted_data, pattern_bytes, known_mappings)
        
        if positions:
            print(f"Found {desc}: {len(positions)} potential positions")
            
            # For each position, try to extract mappings
            for pos in positions[:3]:  # Check first 3 occurrences
                if pos + len(pattern_bytes) <= len(encrypted_data):
                    encrypted_seq = encrypted_data[pos:pos+len(pattern_bytes)]
                    
                    # Extract character mappings
                    for i, (enc_byte, exp_byte) in enumerate(zip(encrypted_seq, pattern_bytes)):
                        if enc_byte not in known_mappings and exp_byte not in new_mappings.values():
                            new_mappings[enc_byte] = exp_byte
                            print(f"  New mapping: 0x{enc_byte:02x} -> '{chr(exp_byte)}' (0x{exp_byte:02x})")
    
    return new_mappings

def find_pattern_positions(encrypted_data, expected_pattern, known_mappings):
    """Find positions where a pattern might occur using partial mappings"""
    
    positions = []
    
    # Convert expected pattern using known reverse mappings
    reverse_mappings = {v: k for k, v in known_mappings.items()}
    
    # Try to find the pattern by looking for partial matches
    for i in range(len(encrypted_data) - len(expected_pattern) + 1):
        encrypted_seq = encrypted_data[i:i+len(expected_pattern)]
        
        # Count how many bytes match our known mappings
        matches = 0
        for enc_byte, exp_byte in zip(encrypted_seq, expected_pattern):
            if exp_byte in reverse_mappings and reverse_mappings[exp_byte] == enc_byte:
                matches += 1
        
        # If we have a good partial match, consider it a candidate
        if matches >= len(expected_pattern) // 2:  # At least half matches
            positions.append(i)
    
    return positions

def analyze_common_words(encrypted_data, known_mappings):
    """Analyze common English words that might appear in the data"""
    
    print("\n--- Common Word Analysis ---")
    
    new_mappings = {}
    
    # Common words we expect to find in API requests
    common_words = [
        b'model', b'text', b'question', b'session', b'user', b'data',
        b'true', b'false', b'null', b'request', b'response', b'chat',
        b'message', b'content', b'params', b'value', b'string', b'number'
    ]
    
    # Use frequency analysis to find these words
    for word in common_words:
        candidates = find_word_candidates(encrypted_data, word, known_mappings)
        
        if candidates:
            print(f"Found candidates for '{word.decode()}': {len(candidates)} positions")
            
            # Use the most likely candidate
            best_candidate = candidates[0]  # First candidate
            pos, encrypted_seq = best_candidate
            
            # Extract mappings
            for enc_byte, exp_byte in zip(encrypted_seq, word):
                if enc_byte not in known_mappings and exp_byte not in new_mappings.values():
                    new_mappings[enc_byte] = exp_byte
                    print(f"  New mapping: 0x{enc_byte:02x} -> '{chr(exp_byte)}' (0x{exp_byte:02x})")
    
    return new_mappings

def find_word_candidates(encrypted_data, word, known_mappings):
    """Find potential positions where a word might be encrypted"""
    
    candidates = []
    reverse_mappings = {v: k for k, v in known_mappings.items()}
    
    for i in range(len(encrypted_data) - len(word) + 1):
        encrypted_seq = encrypted_data[i:i+len(word)]
        
        # Count known character matches
        matches = 0
        for enc_byte, exp_byte in zip(encrypted_seq, word):
            if exp_byte in reverse_mappings and reverse_mappings[exp_byte] == enc_byte:
                matches += 1
        
        # If we have some matches, it's a candidate
        if matches >= 2:  # At least 2 characters match
            candidates.append((i, encrypted_seq))
    
    # Sort by number of matches (best first)
    candidates.sort(key=lambda x: sum(1 for enc, exp in zip(x[1], word) 
                                    if exp in reverse_mappings and reverse_mappings[exp] == enc),
                    reverse=True)
    
    return candidates[:5]  # Top 5 candidates

def analyze_character_frequencies(encrypted_data, known_mappings):
    """Use character frequency analysis to find likely mappings"""
    
    print("\n--- Character Frequency Analysis ---")
    
    new_mappings = {}
    
    # Count byte frequencies in encrypted data
    byte_freq = Counter(encrypted_data)
    
    # Expected character frequencies in English JSON (approximate)
    english_freq_order = [
        ord(' '), ord('e'), ord('t'), ord('a'), ord('o'), ord('i'), ord('n'), ord('s'),
        ord('h'), ord('r'), ord('"'), ord(','), ord(':'), ord('{'), ord('}'), ord('l'),
        ord('u'), ord('d'), ord('c'), ord('m'), ord('w'), ord('f'), ord('g'), ord('y'),
        ord('p'), ord('b'), ord('v'), ord('k'), ord('j'), ord('x'), ord('q'), ord('z')
    ]
    
    # Get most frequent encrypted bytes
    most_frequent_encrypted = [byte_val for byte_val, count in byte_freq.most_common(50)]
    
    # Remove bytes we already know
    unknown_frequent = [b for b in most_frequent_encrypted if b not in known_mappings]
    
    print(f"Top unknown frequent bytes: {[hex(b) for b in unknown_frequent[:10]]}")
    
    # Map most frequent unknown bytes to most frequent expected characters
    used_chars = set(known_mappings.values())
    available_chars = [c for c in english_freq_order if c not in used_chars]
    
    for i, enc_byte in enumerate(unknown_frequent[:len(available_chars)]):
        if i < len(available_chars):
            exp_char = available_chars[i]
            new_mappings[enc_byte] = exp_char
            print(f"  Frequency mapping: 0x{enc_byte:02x} -> '{chr(exp_char)}' (0x{exp_char:02x})")
    
    return new_mappings

def analyze_pattern_repetitions(encrypted_data, known_mappings):
    """Analyze repeated patterns to find more mappings"""
    
    print("\n--- Pattern Repetition Analysis ---")
    
    new_mappings = {}
    
    # Find repeated sequences of different lengths
    for seq_len in [2, 3, 4]:
        print(f"  Analyzing {seq_len}-byte sequences...")
        
        seq_counts = defaultdict(int)
        seq_positions = defaultdict(list)
        
        for i in range(len(encrypted_data) - seq_len + 1):
            seq = encrypted_data[i:i+seq_len]
            seq_counts[seq] += 1
            seq_positions[seq].append(i)
        
        # Find most common sequences
        common_seqs = sorted(seq_counts.items(), key=lambda x: x[1], reverse=True)[:10]
        
        for seq, count in common_seqs:
            if count >= 3:  # Appears at least 3 times
                # Try to decode this sequence using known mappings
                decoded = try_decode_sequence(seq, known_mappings)
                if decoded and all(32 <= c <= 126 for c in decoded):
                    decoded_str = ''.join(chr(c) for c in decoded)
                    print(f"    Sequence {seq.hex()} ({count}x) -> '{decoded_str}'")
                    
                    # If it's a meaningful sequence, extract mappings
                    if is_meaningful_sequence(decoded_str):
                        for enc_byte, dec_byte in zip(seq, decoded):
                            if enc_byte not in known_mappings and dec_byte not in new_mappings.values():
                                new_mappings[enc_byte] = dec_byte
    
    return new_mappings

def try_decode_sequence(seq, known_mappings):
    """Try to decode a sequence using known mappings"""
    
    decoded = []
    for byte_val in seq:
        if byte_val in known_mappings:
            decoded.append(known_mappings[byte_val])
        else:
            return None  # Can't decode if we don't know all bytes
    
    return decoded

def is_meaningful_sequence(text):
    """Check if a decoded sequence looks meaningful"""
    
    # Common meaningful sequences in JSON
    meaningful_patterns = [
        '": ', '", ', '{"', '"}', '": "', 'el"', 'on"', 'er"', 'ed"',
        'ing', 'tion', 'ment', 'ness', 'able', 'ful'
    ]
    
    return any(pattern in text for pattern in meaningful_patterns)

def create_full_mapping_table(mappings):
    """Create a full 256-byte mapping table"""
    
    print(f"\n--- Creating Full Mapping Table ---")
    
    # Start with identity mapping
    table = bytearray(range(256))
    
    # Apply known mappings
    for src, dst in mappings.items():
        table[src] = dst
    
    print(f"Applied {len(mappings)} known mappings")
    
    # Save the table
    with open('comprehensive_mapping.bin', 'wb') as f:
        f.write(table)
    
    print("Saved to comprehensive_mapping.bin")
    
    return bytes(table)

def test_comprehensive_mapping(mapping_table, encrypted_data):
    """Test the comprehensive mapping table"""
    
    print("\n--- Testing Comprehensive Mapping ---")
    
    try:
        # Apply mapping
        decrypted = bytes(mapping_table[b] for b in encrypted_data)
        
        # Try to decode as text
        decrypted_text = decrypted.decode('utf-8', errors='ignore')
        
        # Count JSON-like patterns
        json_patterns = ['"model"', '":', '{"', '"}', ',"', '":"', '"questionText"']
        pattern_counts = {pattern: decrypted_text.count(pattern) for pattern in json_patterns}
        
        total_patterns = sum(pattern_counts.values())
        print(f"JSON patterns found: {total_patterns}")
        for pattern, count in pattern_counts.items():
            if count > 0:
                print(f"  {pattern}: {count} times")
        
        # Check for valid JSON structure
        if total_patterns > 10:  # Good amount of JSON patterns
            print("✅ Strong JSON structure detected!")
            
            # Save result
            with open('comprehensive_decryption.txt', 'w', encoding='utf-8', errors='ignore') as f:
                f.write(decrypted_text)
            
            with open('comprehensive_decryption.json', 'w', encoding='utf-8', errors='ignore') as f:
                # Try to extract JSON
                json_start = decrypted_text.find('{"')
                if json_start != -1:
                    json_end = decrypted_text.rfind('"}') + 2
                    if json_end > json_start:
                        potential_json = decrypted_text[json_start:json_end]
                        f.write(potential_json)
            
            print("Results saved to comprehensive_decryption.txt and .json")
            
            # Show preview
            preview = decrypted_text[:300]
            print(f"Preview: {preview}")
            
            return True
        else:
            print("❌ Insufficient JSON structure found")
            
    except Exception as e:
        print(f"Error testing mapping: {e}")
    
    return False

if __name__ == "__main__":
    print("Comprehensive Mapping Table Extractor")
    print("=" * 50)
    
    if not os.path.exists('decoded_result.bin'):
        print("❌ decoded_result.bin not found")
        exit(1)
    
    # Build comprehensive mapping
    mapping_table, all_mappings = build_comprehensive_mapping()
    
    print(f"\n{'=' * 50}")
    print(f"SUCCESS! Discovered {len(all_mappings)} total mappings")
    print("Final mapping table saved as comprehensive_mapping.bin")
    
    # Show all discovered mappings
    print("\nAll discovered mappings:")
    for src, dst in sorted(all_mappings.items()):
        char_repr = f"'{chr(dst)}'" if 32 <= dst <= 126 else f"0x{dst:02x}"
        print(f"  0x{src:02x} -> {char_repr}")
    
    print("\nCheck comprehensive_decryption.txt for the final result!")